# 主爬虫逻辑
import asyncio
from typing import List, Optional
from playwright.async_api import async_playwright


class WebCrawler:
    def __init__(self, headless: bool = False, timeout: int = 30000, proxy: Optional[dict] = None):
        """
        初始化 WebCrawler。
        :param headless: 是否以无头模式运行浏览器
        :param timeout: 页面加载超时时间（毫秒）
        :param proxy: 代理配置（如 {"server": "http://your-proxy-server:port"}）
        """
        self.headless = headless
        self.timeout = timeout
        self.proxy = proxy

    async def run_playwright(self, url: str, start_page: int, end_page: int) -> List[str]:
        """
        使用 Playwright 直接爬取网页内容。
        :param url: 目标网页的 URL
        :param start_page: 起始页
        :param end_page: 结束页
        :return: 包含每页 HTML 内容的列表
        """
        async with async_playwright() as p:
            # 启动浏览器，配置代理
            browser = await p.chromium.launch(headless=self.headless, proxy=self.proxy)
            # 创建上下文并设置 User-Agent
            context = await browser.new_context(
                user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
            )
            page = await context.new_page()
            all_html_contents = []

            try:
                # 导航到目标网页
                await self._navigate_to_page(page, url)
                # 爬取指定页数范围的数据
                all_html_contents = await self._crawl_pages(page, start_page, end_page)
            except Exception as e:
                print(f"爬取过程中出错: {e}")
            finally:
                await context.close()
                await browser.close()

            return all_html_contents

    async def _navigate_to_page(self, page, url: str):
        """
        导航到目标网页。
        :param page: Playwright 页面对象
        :param url: 目标网页的 URL
        """
        try:
            await page.goto(url, wait_until='networkidle', timeout=self.timeout)
        except asyncio.TimeoutError:
            print("页面加载超时，继续处理已加载的内容。")

    async def _crawl_pages(self, page, start_page: int, end_page: int) -> List[str]:
        """
        爬取指定页数范围的数据。
        :param page: Playwright 页面对象
        :param start_page: 起始页
        :param end_page: 结束页
        :return: 包含每页 HTML 内容的列表
        """
        all_html_contents = []
        for i in range(start_page, end_page + 1):
            print(f"正在爬取第 {i} 页数据...")
            # 获取当前页面的 HTML 内容
            page_source = await page.content()
            all_html_contents.append(page_source)

            # 向下滚动页面
            await self._scroll_page(page)

            # 移除遮罩层（如果存在）
            await self._remove_mask(page)

            # 点击“下一页”按钮
            next_button = await page.query_selector('a.n')
            if next_button:
                await next_button.click()
                # 等待下一页内容加载
                await page.wait_for_selector('div.result', timeout=30000)
            else:
                print("没有更多页面了。")
                break

        return all_html_contents

    async def _scroll_page(self, page):
        """
        向下滚动页面。
        :param page: Playwright 页面对象
        """
        await page.evaluate("window.scrollTo(0, document.body.scrollHeight);")
        await asyncio.sleep(2)  # 等待页面加载新内容

    async def _remove_mask(self, page):
        """
        移除遮罩层。
        :param page: Playwright 页面对象
        """
        await page.evaluate('document.getElementById("_mask")?.remove()')

    async def _go_to_next_page(self, page) -> bool:
        """
        点击“下一页”按钮。
        :param page: Playwright 页面对象
        :return: 是否成功跳转到下一页
        """
        # 查找下一页按钮（根据实际网页调整选择器）
        next_button = await page.query_selector('a.next')  # 例如：豆瓣的下页按钮可能是 a.next
        if next_button:
            # 点击下一页按钮
            await next_button.click()
            # 等待新页面内容加载完成
            try:
                await page.wait_for_selector('div.result', timeout=30000)  # 等待搜索结果加载
                return True
            except Exception as e:
                print(f"等待下一页内容加载失败: {e}")
                return False
        return False
